#include "SphericalHarmonic.h"
#include "Utils.h"
#include <cmath>
#include <iostream>

// 使用Utils命名空间
using namespace Utils;

SphericalHarmonic::SphericalHarmonic(int maxDegree) : maxDegree_(maxDegree) {
    // 初始化
}

SphericalHarmonic::~SphericalHarmonic() {
    // 清理资源
}

void SphericalHarmonic::setVTECObservations(const std::vector<VTECObservation>& vtecObs) {
    vtecObservations_ = vtecObs;
}

void SphericalHarmonic::fitModel() {
    // 使用libtorch实现球谐函数拟合
    
    // 1. 构建设计矩阵
    torch::Tensor designMatrix = buildDesignMatrix();
    
    // 2. 准备观测向量
    auto [vtecTensor, weightsTensor] = observationsToTensor();
    
    // 3. 使用加权最小二乘法求解球谐系数
    // 应用平方根权重以便在矩阵运算中正确应用权重
    torch::Tensor sqrtWeights = torch::sqrt(weightsTensor);
    
    // 对设计矩阵和观测向量应用权重
    torch::Tensor weightedDesignMatrix = designMatrix.clone();
    torch::Tensor weightedVtecTensor = vtecTensor.clone();
    
    for (int i = 0; i < designMatrix.size(0); i++) {
        weightedDesignMatrix[i] = designMatrix[i] * sqrtWeights[i];
        weightedVtecTensor[i] = vtecTensor[i] * sqrtWeights[i];
    }
    
    // 使用QR分解或SVD求解，处理可能的病态问题
    try {
        // 尝试使用常规求解
        auto HTWH = torch::matmul(weightedDesignMatrix.transpose(0, 1), weightedDesignMatrix);
        auto HTWy = torch::matmul(weightedDesignMatrix.transpose(0, 1), weightedVtecTensor);
        coefficients_ = torch::linalg::solve(HTWH, HTWy, true);
    } catch (const c10::Error& e) {
        std::cout << "常规求解失败，使用SVD求解: " << e.what() << std::endl;
        // 使用SVD进行求解
        auto [U, S, Vh] = torch::linalg::svd(
            weightedDesignMatrix,
            true,  // full_matrices
            std::nullopt  // default driver
        );
        
        // 过滤掉过小的奇异值，避免数值不稳定
        double threshold = 1e-10 * torch::max(S).item().toDouble();
        torch::Tensor Sinv = torch::zeros_like(S);
        for (int i = 0; i < S.size(0); i++) {
            if (S[i].item().toDouble() > threshold) {
                Sinv[i] = 1.0 / S[i].item().toDouble();
            }
        }
        
        // 计算解：V * Sinv * U.T * y
        auto V = Vh.transpose(-2, -1);
        auto USinv = U * Sinv.unsqueeze(0);
        auto USinvUT = torch::matmul(USinv, U.transpose(0, 1));
        coefficients_ = torch::matmul(torch::matmul(V, USinvUT), weightedVtecTensor);
    }
    
    std::cout << "球谐函数系数数量: " << coefficients_.size(0) << std::endl;
}

torch::Tensor SphericalHarmonic::buildDesignMatrix() const {
    // 构建球谐函数设计矩阵
    int numObs = vtecObservations_.size();
    int numCoefficients = (maxDegree_ + 1) * (maxDegree_ + 1);
    
    auto options = torch::TensorOptions().dtype(torch::kDouble);
    torch::Tensor designMatrix = torch::zeros({numObs, numCoefficients}, options);
    
    int row = 0;
    for (const auto& obs : vtecObservations_) {
        // 转换为余纬度和弧度
        double colatitude = (90.0 - obs.latitude) * DEG2RAD;
        double longitude = obs.longitude * DEG2RAD;
        
        int col = 0;
        for (int n = 0; n <= maxDegree_; ++n) {
            for (int m = 0; m <= n; ++m) {
                // 计算勒让德函数
                double pnm = normalizedLegendre(n, m, std::cos(colatitude));
                
                // 余弦项
                designMatrix[row][col] = pnm * std::cos(m * longitude);
                col++;
                
                // 正弦项 (m > 0时)
                if (m > 0) {
                    designMatrix[row][col] = pnm * std::sin(m * longitude);
                    col++;
                }
            }
        }
        row++;
    }
    
    return designMatrix;
}

std::tuple<torch::Tensor, torch::Tensor> SphericalHarmonic::observationsToTensor() const {
    int numObs = vtecObservations_.size();
    
    auto options = torch::TensorOptions().dtype(torch::kDouble);
    torch::Tensor vtecTensor = torch::zeros({numObs}, options);
    torch::Tensor weightsTensor = torch::zeros({numObs}, options);
    
    for (int i = 0; i < numObs; ++i) {
        vtecTensor[i] = vtecObservations_[i].vtec;
        
        // 修改为直接使用高度角的余弦值作为权重
        // 高度角越高(接近90度)，cos值越小，权重越大
        double elevation = vtecObservations_[i].elevation * DEG2RAD;
        // 使用余弦的倒数或者sin(elevation)作为权重
        // 这里使用sin(elevation)，等同于cos(zenith)
        weightsTensor[i] = std::sin(elevation);
    }
    
    return {vtecTensor, weightsTensor};
}

double SphericalHarmonic::predictVTEC(double latitude, double longitude, double time) {
    // Implementation
    // Mark time as used to suppress warning
    (void) time;
    
    // Convert to colatitude (0 at north pole, π at south pole)
    double colatitude = (90.0 - latitude) * DEG2RAD;
    
    // Convert to longitude in radians
    double lonRad = longitude * DEG2RAD;
    
    double vtec = 0.0;
    int idx = 0;
    
    for (int n = 0; n <= maxDegree_; n++) {
        for (int m = 0; m <= n; m++) {
            double Pnm = Utils::normalizedLegendre(n, m, std::cos(colatitude));
            
            if (m == 0) {
                vtec += coefficients_[idx].item().toDouble() * Pnm;
                idx++;
            } else {
                vtec += coefficients_[idx].item().toDouble() * Pnm * std::cos(m * lonRad);
                idx++;
                vtec += coefficients_[idx].item().toDouble() * Pnm * std::sin(m * lonRad);
                idx++;
            }
        }
    }
    
    return vtec;
}

torch::Tensor SphericalHarmonic::getCoefficients() const {
    return coefficients_;
}

torch::Tensor SphericalHarmonic::generateGlobalVTECGrid(int latRes, int lonRes) const {
    int numLat = 180 / latRes + 1;
    int numLon = 360 / lonRes;
    
    auto options = torch::TensorOptions().dtype(torch::kFloat);
    torch::Tensor vtecGrid = torch::zeros({numLat, numLon}, options);
    
    // For const correctness, use a non-const version of this object to call the non-const method
    // This is safe because we're not modifying the object state
    const SphericalHarmonic* constThis = this;
    SphericalHarmonic* mutableThis = const_cast<SphericalHarmonic*>(constThis);
    
    for (int i = 0; i < numLat; ++i) {
        double lat = 90.0 - i * latRes;
        
        for (int j = 0; j < numLon; ++j) {
            double lon = j * lonRes;
            
            // Use the mutable version to call the non-const method
            vtecGrid[i][j] = mutableThis->predictVTEC(lat, lon, 0.0);
        }
    }
    
    return vtecGrid;
}

std::vector<double> SphericalHarmonic::getCoefficientsVector() const {
    std::vector<double> coeffsVec;
    
    // 如果系数尚未计算，返回空向量
    if (!coefficients_.defined()) {
        return coeffsVec;
    }
    
    // 将torch::Tensor转换为std::vector<double>
    auto accessor = coefficients_.accessor<float, 1>();
    coeffsVec.reserve(accessor.size(0));
    
    for (int i = 0; i < accessor.size(0); ++i) {
        coeffsVec.push_back(accessor[i]);
    }
    
    return coeffsVec;
}

void SphericalHarmonic::setCoefficients(const std::vector<double>& coeffs) {
    // 创建torch::Tensor并填充系数
    coefficients_ = torch::zeros({static_cast<int64_t>(coeffs.size())});
    auto accessor = coefficients_.accessor<float, 1>();
    
    for (size_t i = 0; i < coeffs.size(); ++i) {
        accessor[i] = static_cast<float>(coeffs[i]);
    }
}

double SphericalHarmonic::evaluateVTEC(double latitude, double longitude) const {
    if (!coefficients_.defined()) {
        std::cerr << "球谐系数尚未设置，无法评估VTEC" << std::endl;
        return 0.0;
    }
    
    // 转换为弧度
    double lat_rad = latitude * M_PI / 180.0;
    double lon_rad = longitude * M_PI / 180.0;
    
    // 计算余纬
    double coLat = M_PI / 2.0 - lat_rad;
    
    // 计算球谐函数值
    double vtec = 0.0;
    int idx = 0;
    
    for (int n = 0; n <= maxDegree_; ++n) {
        for (int m = 0; m <= n; ++m) {
            // 计算勒让德多项式
            double pnm = Utils::normalizedLegendre(n, m, std::cos(coLat));
            
            // 计算球谐函数
            if (m == 0) {
                vtec += coefficients_[idx].item().toDouble() * pnm;
                idx++;
            } else {
                // 余弦项
                vtec += coefficients_[idx].item().toDouble() * pnm * std::cos(m * lon_rad);
                idx++;
                
                // 正弦项
                vtec += coefficients_[idx].item().toDouble() * pnm * std::sin(m * lon_rad);
                idx++;
            }
        }
    }
    
    return vtec;
}

// 计算勒让德多项式值
double SphericalHarmonic::calculateLegendreValue(int n, int m, double x) const {
    // 特殊情况处理
    if (m > n) return 0.0;
    if (n == 0 && m == 0) return 1.0;
    
    // 递归计算
    if (n == m) {
        if (m == 0) return 1.0;
        return (2.0 * m - 1.0) * std::sqrt(1.0 - x * x) * calculateLegendreValue(m - 1, m - 1, x);
    }
    
    if (n == m + 1) {
        return (2.0 * m + 1.0) * x * calculateLegendreValue(m, m, x);
    }
    
    // 一般情况
    return ((2.0 * n - 1.0) * x * calculateLegendreValue(n - 1, m, x) - 
            (n + m - 1.0) * calculateLegendreValue(n - 2, m, x)) / (n - m);
} 